/*
 * This file is a part of the SWFer project.
 *
 * Copyright (C) 2012 Michael Bradshaw <mjbshaw@gmail.com>
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program; if not, write to the Free Software Foundation, Inc., 59 Temple
 * Place, Suite 330, Boston, MA 02111-1307 USA
 *
 */

#include "NewSimpleSound.hpp"
#include "ui_NewSimpleSound.h"

#include "SwfMainDisplay.hpp"

#include <iostream>

#include <swfer/SimpleSwf.hpp>
#include <swfer/SimpleSound.hpp>

#include <QFileDialog>
#include <QThread>

extern "C"
{
#include <libavformat/avformat.h>
#include <libswresample/swresample.h>
}

NewSimpleSound::NewSimpleSound(QWidget* parent) : QDialog(parent),
    ui(new Ui::NewSimpleSound)
{
    ui->setupUi(this);
	
	connect(ui->selectFileButton, SIGNAL(pressed()), this, SLOT(selectFile()));
	connect(ui->saveAsButton, SIGNAL(pressed()), this, SLOT(saveAs()));
	connect(ui->cancelButton, SIGNAL(pressed()), this, SLOT(close()));
	
	//swfer::SimpleSwf tmp1("../samples/UntitledTest.swf");
	//swfer::SimpleSwf tmp2("../samples/zhacked.swf");
	//swfer::SimpleSwf tmp3("../samples/nphacked.swf");
	//swfer::SimpleSwf tmp4("../samples/TEST_WITH_PACKAGE.swf");
	//swfer::SimpleSwf tmp5("../samples/TEST_NO_PACKAGE.swf");
	//swfer::SimpleSwf tmp6("../samples/sineloop.swf");
	//swfer::SimpleSwf tmp7("../samples/np.swf");
	//swfer::SimpleSwf tmp8("../samples/ha.swf");
}

NewSimpleSound::~NewSimpleSound()
{
    delete ui;
}

void NewSimpleSound::updateProgress(int percentage)
{
	ui->progressBar->setValue(percentage);
	ui->progressBar->update();
}

#include <fstream>

struct Encoder
{
	AVCodec* codec;
	AVCodecContext* ctx;
	AVFrame* frame;
	int bufferSize;
	uint16_t* samples;
	std::size_t samplesInFrame;

	int numSamplesEncoded;

	std::ofstream outFile;

	Encoder(int sr, int aq, int vbrq)
	{
		outFile.open("../zzzout.mp3", std::ios::binary);
		
		samplesInFrame = 0;
		numSamplesEncoded = 0;

		codec = avcodec_find_encoder(AV_CODEC_ID_MP3);
		if (!codec)
		{
			// error
			std::cerr << "Could not find mp3 codec" << std::endl;
		}
		
		ctx = avcodec_alloc_context3(codec);

		ctx->flags |= CODEC_FLAG_QSCALE; // VBR
		ctx->global_quality = vbrq * FF_QP2LAMBDA; // [0, 9] for lame encoder
		ctx->compression_level = aq; // [0, 9] for lame encoder
		//ctx->bit_rate = 64000; // for CBR

		ctx->sample_fmt = AV_SAMPLE_FMT_S16;
		// TODO: assert ctx->sample_fmt is in codec->sample_fmts

		ctx->sample_rate = sr;
		// TODO: check codec->sample_rates
		ctx->channel_layout = AV_CH_LAYOUT_STEREO;
		// TODO: check codec->channel_layouts
		ctx->channels = av_get_channel_layout_nb_channels(ctx->channel_layout);

		if (avcodec_open2(ctx, codec, NULL) < 0)
		{
			// error
			std::cerr << "Couldn't open codec" << std::endl;
		}

		frame = avcodec_alloc_frame();
		if (!frame)
		{
			// error
			std::cerr << "Couldn't allocate frame" << std::endl;
		}

		frame->nb_samples = ctx->frame_size;
		frame->format = ctx->sample_fmt;
		frame->channel_layout = ctx->channel_layout;

		bufferSize = av_samples_get_buffer_size(NULL, ctx->channels,
			ctx->frame_size, ctx->sample_fmt, 0);

		samples = static_cast<uint16_t*>(av_malloc(bufferSize));
		if (!samples)
		{
			// error
			std::cerr << "Couldn't allocate samples buffer" << std::endl;
		}

		int ret = avcodec_fill_audio_frame(frame, ctx->channels, ctx->sample_fmt,
			(const uint8_t*)samples, bufferSize, 0);
		if (ret < 0)
		{
			// error
			std::cerr << "Couldnt set up audio frame" << std::endl;
		}
	}

	void processSamples(const char* buffer, int nbSamples)
	{
		for (int i = 0; i < nbSamples; ++i)
		{
			// 2 because 2 == sizeof(uint16_t)
			samples[2 * samplesInFrame] = reinterpret_cast<const uint16_t*>(buffer)[2 * i];

			for (int k = 1; k < ctx->channels; ++k)
			{
				samples[2 * samplesInFrame + k] = reinterpret_cast<const uint16_t*>(buffer)[2 * i + k];
			}

			++samplesInFrame;

			if (samplesInFrame == ctx->frame_size)
			{
				encode();
			}
		}
	}

	void finish()
	{
		if (samplesInFrame > 0)
		{
			frame->nb_samples = samplesInFrame;
			encode();
		}

		if (codec->capabilities & CODEC_CAP_DELAY)
		{
			AVPacket pkt;
			av_init_packet(&pkt);
			pkt.data = NULL;
			pkt.size = 0;

			int gotPacket = 0;
			while (avcodec_encode_audio2(ctx, &pkt, NULL, &gotPacket) == 0 && gotPacket)
			{
				outFile.write((const char*)pkt.data, pkt.size);
				for (int i = 0; i < pkt.size; ++i)
				{
					data.push_back(pkt.data[i]);
				}

				av_free_packet(&pkt);
			}
		}

		outFile.close();
	}

	void encode()
	{
		AVPacket pkt;
		av_init_packet(&pkt);
		pkt.data = NULL;
		pkt.size = 0;

		int gotPacket = 0;
		int ret = avcodec_encode_audio2(ctx, &pkt, frame, &gotPacket);
		if (ret < 0)
		{
			// error
			std::cerr << "Error encoding audio frame" << std::endl;
		}
		
		if (gotPacket)
		{
			outFile.write((const char*)pkt.data, pkt.size);
			for (int i = 0; i < pkt.size; ++i)
			{
				data.push_back(pkt.data[i]);
			}

			av_free_packet(&pkt);
		}
		
		numSamplesEncoded += samplesInFrame;
		samplesInFrame = 0;
	}

	std::vector<char> data;
};

#include <fstream>

bool NewSimpleSound::encode(const QString& sourcePath, swfer::SimpleSound& ss, swfer::SampleRate sr, int aq, int vbrq)
{
	Encoder encoder((sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0), aq, vbrq);

    AVFrame* frame = avcodec_alloc_frame();
    if (!frame)
    {
        std::cerr << "Error allocating the frame" << std::endl;
        return false;
    }

    AVFormatContext* formatContext = NULL;
	if (avformat_open_input(&formatContext, sourcePath.toStdString().c_str(), NULL, NULL) != 0)
    {
        av_free(frame);
        std::cerr << "Error opening the file" << std::endl;
        return false;
    }

    if (avformat_find_stream_info(formatContext, NULL) < 0)
    {
        av_free(frame);
        avformat_close_input(&formatContext);
        std::cerr << "Error finding the stream info" << std::endl;
        return false;
    }

    AVStream* audioStream = NULL;
    for (unsigned int i = 0; i < formatContext->nb_streams; ++i)
    {
        if (formatContext->streams[i]->codec->codec_type == AVMEDIA_TYPE_AUDIO)
        {
            audioStream = formatContext->streams[i];
            break;
        }
    }

    if (audioStream == NULL)
    {
        av_free(frame);
        avformat_close_input(&formatContext);
        std::cerr << "Could not find any audio stream in the file" << std::endl;
        return false;
    }

    AVCodecContext* codecContext = audioStream->codec;

    codecContext->codec = avcodec_find_decoder(codecContext->codec_id);
    if (codecContext->codec == NULL)
    {
        av_free(frame);
        avformat_close_input(&formatContext);
        std::cerr << "Couldn't find a proper decoder" << std::endl;
        return false;
    }
    else if (avcodec_open2(codecContext, codecContext->codec, NULL) != 0)
    {
        av_free(frame);
        avformat_close_input(&formatContext);
        std::cerr << "Couldn't open the context with the decoder" << std::endl;
        return false;
    }

    //std::cout << "This stream has " << codecContext->channels << " channels and a sample rate of " << codecContext->sample_rate << "Hz" << std::endl;
    //std::cout << "The data is in the format " << av_get_sample_fmt_name(codecContext->sample_fmt) << std::endl;

    int64_t outChannelLayout = AV_CH_LAYOUT_STEREO;
    AVSampleFormat outSampleFormat = AV_SAMPLE_FMT_S16;
	int outSampleRate = (sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0);

    // Note that AVCodecContext::channel_layout may or may not be set by libavcodec. Because of this,
    // we won't use it, and will instead try to guess the layout from the number of channels.
    SwrContext* swrContext = swr_alloc_set_opts(NULL,
                                                outChannelLayout,
                                                outSampleFormat,
                                                outSampleRate,
                                                av_get_default_channel_layout(codecContext->channels),
                                                codecContext->sample_fmt,
                                                codecContext->sample_rate,
                                                0,
                                                NULL);

    if (swrContext == NULL)
    {
        av_free(frame);
        avcodec_close(codecContext);
        avformat_close_input(&formatContext);
        std::cerr << "Couldn't create the SwrContext" << std::endl;
        return false;
    }

    if (swr_init(swrContext) != 0)
    {
        av_free(frame);
        avcodec_close(codecContext);
        avformat_close_input(&formatContext);
        swr_free(&swrContext);
        std::cerr << "Couldn't initialize the SwrContext" << std::endl;
        return false;
    }

    AVPacket packet;
    av_init_packet(&packet);

    // Read the packets in a loop
	int totalRead = 0;
    while (av_read_frame(formatContext, &packet) == 0)
    {
        if (packet.stream_index == audioStream->index)
        {
            // Create a copy of the packet that we can modify and manipulate
            AVPacket decodingPacket = packet;

            while (decodingPacket.size > 0)
            {
                // Try to decode the packet into a frame
                int frameFinished = 0;
                int result = avcodec_decode_audio4(codecContext, frame, &frameFinished, &decodingPacket);

                if (result < 0 || frameFinished == 0)
                {
                    break;
                }

				totalRead += frame->nb_samples;

				std::vector<unsigned char> buffer(2 * (sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0) * 2);

				unsigned char* outPtrs[SWR_CH_MAX] = {NULL};
				outPtrs[0] = &buffer[0];

				int numSamplesOut = swr_convert(swrContext,
												outPtrs,
												(sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0),
												(const uint8_t**)frame->data,
												frame->nb_samples);

				encoder.processSamples((const char*)&buffer[0], numSamplesOut);

                decodingPacket.size -= result;
                decodingPacket.data += result;

				swfer::uint64_t progress = totalRead;
				progress *= 100;
				progress /= audioStream->codec->sample_rate;

				progress *= AV_TIME_BASE;
				progress /= formatContext->duration;

				ui->progressBar->setValue((int)progress);
				QApplication::processEvents();
            }
        }

        // You *must* call av_free_packet() after each call to av_read_frame() or else you could leak memory
        av_free_packet(&packet);
    }

    // Some codecs will cause frames to be buffered up in the decoding process. If the CODEC_CAP_DELAY flag
    // is set, there can be buffered up frames that need to be flushed, so we'll do that
    if (codecContext->codec->capabilities & CODEC_CAP_DELAY)
    {
        av_init_packet(&packet);
        // Decode all the remaining frames in the buffer, until the end is reached
        int frameFinished = 0;
        while (avcodec_decode_audio4(codecContext, frame, &frameFinished, &packet) >= 0 && frameFinished)
        {
			std::vector<unsigned char> buffer(2 * (sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0) * 2);

			unsigned char* outPtrs[SWR_CH_MAX] = {NULL};
			outPtrs[0] = &buffer[0];

			int numSamplesOut = swr_convert(swrContext,
											outPtrs,
											(sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0),
											(const uint8_t**)frame->data,
											frame->nb_samples);

			encoder.processSamples((const char*)&buffer[0], numSamplesOut);
        }
    }

    // swr_convert can buffer data up. Passing NULL and 0 into it will flush what's buffered out.
    // You should do this at the very end. Alternatively, if you always want to get every sample
    // out (i.e. you don't want it to buffer), you can do this after each call to swr_convert.
    int numSamplesOut = 0;
    do {
        std::vector<unsigned char> buffer(2 * (sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0) * 2);

        unsigned char* pointers[SWR_CH_MAX] = {NULL};
        pointers[0] = &buffer[0];

        numSamplesOut = swr_convert(swrContext,
                                    pointers,
                                    (sr == swfer::SAMPLE_44100_HZ ? 44100 :
		sr == swfer::SAMPLE_22050_HZ ? 22050 :
		sr == swfer::SAMPLE_11025_HZ ? 11025 :
		0),
                                    NULL,
                                    0);

		encoder.processSamples((const char*)&buffer[0], numSamplesOut);
    } while (numSamplesOut > 0);

	encoder.finish();

	ss.format = swfer::MP3;
	ss.samples = encoder.numSamplesEncoded;
	ss.skip = audioStream->codec->delay + encoder.ctx->delay;
	ss.data.swap(encoder.data);
	ss.size = swfer::SAMPLE_16_BITS;
	ss.layout = swfer::STEREO;
	ss.rate = sr;

    // Clean up!
    av_free(frame);
    avcodec_close(codecContext);
    avformat_close_input(&formatContext);
    swr_free(&swrContext);
	
	ui->progressBar->setValue(100);
	return true;
}

void NewSimpleSound::saveAs()
{
    QString path = QFileDialog::getSaveFileName(this, QString(), QString(), "SWF file (*.swf)");
	if (!path.isEmpty())
	{
		// Note: this is leaking...
		//Transcoder* t = new Transcoder(*this, ui->sourceFileLineEdit->text(), path, ui->classNameLineEdit->text(), (swfer::SampleRate)(ui->sampleRateComboBox->currentIndex() + 1), ui->algorithmicQualitySlider->value(), ui->vbrQualitySlider->value());
		//t->start();

		swfer::SimpleSound ss;
		ss.className = ui->classNameLineEdit->text().toStdString();
		//ss.packageName = ui->packageNameLineEdit->text().toStdString();
		//ss.loop = ui->loopCheckBox->isChecked();
		// +1 since 5512 isn't supported by encoder
		encode(ui->sourceFileLineEdit->text(), ss, (swfer::SampleRate)(ui->sampleRateComboBox->currentIndex() + 1), ui->algorithmicQualitySlider->value(), ui->vbrQualitySlider->value());

		swfer::SimpleSwf swf;
		swf.addSound(ss);
		swf.save(path.toStdString());
	}
}

void NewSimpleSound::selectFile()
{
    QString path = QFileDialog::getOpenFileName(this, QString(), QString(), "Lossless audio file (*.wav *.flac)");
	if (!path.isEmpty())
	{
		ui->sourceFileLineEdit->setText(path);
	}
}